"""
综合量化交易策略融合动量、双均线、均值回归、布林带和成交量五大因子，
通过动态加权整合各子策略信号，突破单一策略局限：动量策略（20%）捕捉趋势惯性，
双均线（25%）判断趋势转折，布林带（20%）识别超买超卖，均值回归（15%）捕捉反转机会，
成交量（20%）验证价格变动有效性，最终生成梯度化交易信号，并通过可视化矩阵直观展示多空强度，
适用于趋势跟踪、反转交易等多种市场环境。
"""
import pandas as pd
import numpy as np
from sqlalchemy import create_engine
import matplotlib.pyplot as plt
from datetime import datetime

# 创建数据库连接引擎
engine = create_engine('mysql+pymysql://root:@127.0.0.1:3306/stocks?charset=utf8mb4&use_unicode=1')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


class ComprehensiveStrategy:
    def __init__(self, lookback_period=20, volume_multiple=2.0,
                 short_window=5, long_window=20,
                 bollinger_window=20, num_std=2,
                 mean_reversion_window=10, z_score_threshold=1.0):
        """
        初始化综合策略参数
        参数:
            lookback_period: 成交量回溯期(天数)
            volume_multiple: 成交量倍数阈值
            short_window: 短期均线窗口
            long_window: 长期均线窗口
            bollinger_window: 布林带窗口
            num_std: 布林带标准差倍数
            mean_reversion_window: 均值回归窗口
            z_score_threshold: Z分数阈值
        """
        self.lookback_period = lookback_period
        self.volume_multiple = volume_multiple
        self.short_window = short_window
        self.long_window = long_window
        self.bollinger_window = bollinger_window
        self.num_std = num_std
        self.mean_reversion_window = mean_reversion_window
        self.z_score_threshold = z_score_threshold

    def get_stock_data(self, ts_code, days):
        """从数据库获取股票数据"""
        query = f"""
            SELECT trade_date, close, vol, pct_chg 
            FROM stocks_daily 
            WHERE ts_code = '{ts_code}' 
            ORDER BY trade_date DESC 
            LIMIT {days}
        """
        df = pd.read_sql(query, con=engine)
        df = df.sort_values('trade_date')
        return df

    def get_latest_price(self, ts_code):
        """获取股票最新价格"""
        try:
            query = f"""
                SELECT close 
                FROM stocks_daily 
                WHERE ts_code = '{ts_code}' 
                ORDER BY trade_date DESC 
                LIMIT 1
            """
            result = pd.read_sql(query, con=engine)
            if len(result) > 0:
                return round(result['close'].iloc[0], 2)
            else:
                return "N/A"
        except Exception as e:
            print(f"获取{ts_code}最新价格失败: {e}")
            return "N/A"

    def momentum_strategy(self, df):
        """动量策略（使用pct_chg字段）"""
        df['momentum'] = df['pct_chg'].rolling(self.lookback_period).apply(
            lambda x: (x / 100 + 1).prod() - 1
        )
        df['momentum_signal'] = np.where(
            df['momentum'] > 0, 1,  # 正动量买入
            np.where(df['momentum'] < 0, -1, 0)  # 负动量卖出
        )
        return df

    def moving_average_crossover(self, df):
        """双均线策略"""
        df['short_ma'] = df['close'].rolling(window=self.short_window).mean()
        df['long_ma'] = df['close'].rolling(window=self.long_window).mean()

        # 生成信号
        df['ma_signal'] = np.where(
            df['short_ma'] > df['long_ma'], 1,  # 金叉买入
            np.where(df['short_ma'] < df['long_ma'], -1, 0)  # 死叉卖出
        )
        return df

    def bollinger_bands(self, df):
        """布林带策略"""
        df['rolling_mean'] = df['close'].rolling(window=self.bollinger_window).mean()
        df['rolling_std'] = df['close'].rolling(window=self.bollinger_window).std()
        df['upper_band'] = df['rolling_mean'] + (df['rolling_std'] * self.num_std)
        df['lower_band'] = df['rolling_mean'] - (df['rolling_std'] * self.num_std)

        # 生成信号
        df['bollinger_signal'] = np.where(
            df['close'] < df['lower_band'], 1,  # 低于下轨买入
            np.where(df['close'] > df['upper_band'], -1, 0)  # 高于上轨卖出
        )
        return df

    def mean_reversion(self, df):
        """均值回归策略"""
        df['mean'] = df['close'].rolling(window=self.mean_reversion_window).mean()
        df['std'] = df['close'].rolling(window=self.mean_reversion_window).std()
        df['z_score'] = (df['close'] - df['mean']) / df['std']

        # 生成信号
        df['mean_reversion_signal'] = np.where(
            df['z_score'] < -self.z_score_threshold, 1,  # 低于均值买入
            np.where(df['z_score'] > self.z_score_threshold, -1, 0)  # 高于均值卖出
        )
        return df

    def volume_strategy(self, df):
        """成交量策略（使用vol和pct_chg字段）"""
        df['vol_ma'] = df['vol'].rolling(window=self.lookback_period).mean()
        df['vol_ratio'] = df['vol'] / df['vol_ma']

        # 生成信号
        df['volume_signal'] = np.where(
            (df['vol_ratio'] >= self.volume_multiple) & (df['pct_chg'] > 0), 1,  # 放量上涨买入
            np.where(
                (df['vol_ratio'] >= self.volume_multiple) & (df['pct_chg'] < 0), -1,  # 放量下跌卖出
                np.where(
                    (df['vol_ratio'] < self.volume_multiple) & (df['pct_chg'] > 0), 0.5,  # 量平上涨持有
                    0  # 其他情况观望
                )
            )
        )
        return df

    def generate_signals(self):
        """生成综合交易信号"""
        # 获取股票列表
        query = "SELECT ts_code, name FROM stocks_info"
        stocks = pd.read_sql(query, con=engine)

        results = []
        for index, stock in stocks.iterrows():
            ts_code = stock['ts_code']
            name = stock['name']
            serial_number = index + 1  # 生成序号，从1开始

            try:
                # 获取最新价格
                latest_price = self.get_latest_price(ts_code)
                
                # 获取足够长的数据来计算所有指标
                days_needed = max(self.lookback_period, self.long_window,
                                  self.bollinger_window, self.mean_reversion_window) + 10
                df = self.get_stock_data(ts_code, days_needed)

                if len(df) < days_needed:
                    results.append([serial_number, name, ts_code, latest_price, "数据不足", "", "", "", "", ""])
                    continue

                # 应用各策略
                df = self.momentum_strategy(df)
                df = self.moving_average_crossover(df)
                df = self.bollinger_bands(df)
                df = self.mean_reversion(df)
                df = self.volume_strategy(df)

                # 获取最新信号
                latest = df.iloc[-1]

                # 综合信号 (简单加权平均)
                combined_signal = (0.2 * latest['momentum_signal'] +
                                   0.25 * latest['ma_signal'] +
                                   0.2 * latest['bollinger_signal'] +
                                   0.15 * latest['mean_reversion_signal'] +
                                   0.2 * latest['volume_signal'])

                # 生成交易建议
                if combined_signal > 0.5:
                    action = "强烈买入"
                elif combined_signal > 0.2:
                    action = "买入"
                elif combined_signal > -0.2:
                    action = "持有"
                elif combined_signal > -0.5:
                    action = "卖出"
                else:
                    action = "强烈卖出"

                results.append([
                    serial_number, name, ts_code, latest_price,
                    f"{action}(综合:{combined_signal:.2f})",
                    "买入" if latest['momentum_signal'] > 0 else "卖出" if latest['momentum_signal'] < 0 else "观望",
                    "买入" if latest['ma_signal'] > 0 else "卖出" if latest['ma_signal'] < 0 else "观望",
                    "买入" if latest['bollinger_signal'] > 0 else "卖出" if latest['bollinger_signal'] < 0 else "观望",
                    "买入" if latest['mean_reversion_signal'] > 0 else "卖出" if latest['mean_reversion_signal'] < 0 else "观望",
                    "买入" if latest['volume_signal'] > 0 else "卖出" if latest['volume_signal'] < 0 else "观望"
                ])

            except Exception as e:
                results.append([serial_number, name, ts_code, "N/A", f"错误:{str(e)}", "", "", "", "", ""])

        # 创建结果DataFrame
        columns = [
            '序号', '股票名称', '股票代码', '最新价格',
            '综合信号', '动量信号', '均线信号', 
            '布林带信号', '均值回归信号', '成交量信号'
        ]
        df_result = pd.DataFrame(results, columns=columns)

        return df_result

    def visualize_signals(self, df_result):
        """可视化信号表格"""
        # 获取信号统计信息
        signal_counts = self.get_signal_statistics_for_display(df_result)
        
        # 动态计算图形尺寸
        num_rows = len(df_result)
        row_height = 0.5
        header_height = 0.8
        footer_space = 0.5
        title_space = 1.0

        # 调整图形高度，不再需要额外空间给统计信息
        fig_height = min(
            title_space + header_height + num_rows * row_height + footer_space,
            30
        )

        # 创建图形 - 增加宽度以容纳右侧统计信息
        fig = plt.figure(figsize=(24, fig_height))  # 增加宽度
        ax = fig.add_subplot(111)
        ax.axis('off')

        # 添加标题和时间戳
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        title = fig.suptitle(
            f'综合量化交易策略信号\n生成时间: {current_time}',
            y=0.96,
            fontsize=14,
            fontweight='bold'
        )

        # 添加策略说明
        strategy_desc = """综合量化交易策略融合动量、双均线、均值回归、布林带和成交量五大因子，
    通过动态加权整合各子策略信号，突破单一策略局限：动量策略（20%）捕捉趋势惯性，
    双均线（25%）判断趋势转折，布林带（20%）识别超买超卖，均值回归（15%）捕捉反转机会，
    成交量（20%）验证价格变动有效性，最终生成梯度化交易信号。"""
        
        # 添加策略说明文本（左侧）
        ax.text(0.02, 0.99, strategy_desc, 
            transform=ax.transAxes, fontsize=9, 
            verticalalignment='top', style='italic',
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

        # 添加信号统计信息（右侧）
        stats_text = "=== 信号分布统计 ===\n"
        total_stocks = len(df_result)
        stats_text += f"总股票数量: {total_stocks}只\n\n"
        
        # 只显示有数据的信号类型
        displayed_signals = []
        for signal_type, count in signal_counts.items():
            if count > 0:
                percentage = (count / total_stocks) * 100
                stats_text += f"{signal_type}: {count}只 ({percentage:.1f}%)\n"
                displayed_signals.append(signal_type)
        
        # 如果没有有效的信号数据，添加提示
        if not displayed_signals:
            stats_text += "暂无有效信号数据"
        
        # 在右侧添加统计信息框
        ax.text(0.55, 0.99, stats_text, 
            transform=ax.transAxes, fontsize=10, 
            verticalalignment='top',
            bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.7))

        # 调整表格位置和大小 - 占据左侧大部分空间
        table_bottom = 0.05
        table_height = 0.80  # 恢复表格高度
        table_width = 0.60   # 表格宽度，为右侧留出空间

        # 创建表格（放在左侧）
        table = ax.table(
            cellText=df_result.values,
            colLabels=df_result.columns,
            loc='center',
            bbox=[0.25, table_bottom, table_width, table_height],
            cellLoc='center'
        )

        # 设置表格样式
        table.auto_set_font_size(False)
        table.set_fontsize(8)  # 稍微减小字体以适应更多内容

        # 设置行高和列宽
        for i in range(len(df_result) + 1):
            for j in range(len(df_result.columns)):
                cell = table[(i, j)]
                cell.set_height(0.06)
                cell.set_edgecolor('lightgray')

        # 设置表头样式
        for j in range(len(df_result.columns)):
            table[(0, j)].set_facecolor('#40466e')
            table[(0, j)].get_text().set_color('white')
            table[(0, j)].set_fontsize(9)

        # 设置信号单元格颜色
        color_map = {
            '强烈买入': '#ff0000',      # 红色
            '买入': '#ff6b6b',         # 浅红色
            '持有': '#f8f9fa',         # 浅灰色
            '卖出': '#51cf66',         # 绿色
            '强烈卖出': '#2b8a3e',     # 深绿色
            '观望': '#339af0',         # 蓝色
            '错误': '#ffd43b',         # 黄色
            '数据不足': '#ced4da'      # 灰色
        }

        # 设置特殊列样式
        for i in range(1, len(df_result) + 1):
            # 序号列样式
            serial_cell = table[(i, 0)]
            serial_cell.set_facecolor('#e9ecef')
            serial_cell.get_text().set_fontweight('bold')
            
            # 价格列样式
            price_cell = table[(i, 3)]
            price_value = df_result.iloc[i-1, 3]
            if price_value != "N/A" and price_value != "数据不足" and price_value != "错误":
                try:
                    float(price_value)  # 检查是否为数字
                    price_cell.set_facecolor('#fff3cd')
                    price_cell.get_text().set_fontweight('bold')
                except (ValueError, TypeError):
                    pass
            
            # 信号列样式
            for j in range(4, len(df_result.columns)):
                signal = str(df_result.iloc[i - 1, j])
                color = 'white'
                for key, val in color_map.items():
                    if key in signal:
                        color = val
                        break

                table[(i, j)].set_facecolor(color)
                # 为深色背景设置白色文字
                if color in ['#ff0000', '#2b8a3e', '#40466e']:
                    table[(i, j)].get_text().set_color('white')
                else:
                    table[(i, j)].get_text().set_color('black')

        # 在右侧添加策略权重说明
        weights_text = "=== 策略权重分配 ===\n"
        weights_text += "• 动量策略: 20%\n"
        weights_text += "• 双均线策略: 25%\n"
        weights_text += "• 布林带策略: 20%\n"
        weights_text += "• 均值回归: 15%\n"
        weights_text += "• 成交量策略: 20%\n\n"
        weights_text += "=== 信号颜色说明 ===\n"
        weights_text += "• 红色: 强烈买入\n"
        weights_text += "• 浅红: 买入\n"
        weights_text += "• 绿色: 卖出\n"
        weights_text += "• 深绿: 强烈卖出\n"
        weights_text += "• 蓝色: 观望\n"
        weights_text += "• 灰色: 持有/数据不足"

        # 在右侧添加权重和颜色说明
        ax.text(0.33, 1.08, weights_text, 
            transform=ax.transAxes, fontsize=9, 
            verticalalignment='top',
            bbox=dict(boxstyle='round', facecolor='lightgreen', alpha=0.7))

        # 调整布局
        plt.tight_layout(rect=[0, 0, 1, 0.96])
        plt.subplots_adjust(top=0.92, bottom=0.05)

        # 添加说明文字
        ax.text(0.02, 0.01, 
            '信号结果仅供参考，投资有风险', 
            transform=ax.transAxes, fontsize=9, style='italic')
        
        plt.show()

    def save_to_excel(self, df_result, filename=None):
        """保存结果到Excel文件"""
        if filename is None:
            current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f'综合策略信号_{current_time}.xlsx'
        
        # 保存时不需要序号列（因为Excel会自动生成行号）
        df_to_save = df_result.copy()
        df_to_save.to_excel(filename, index=False, engine='openpyxl')
        print(f"结果已保存到: {filename}")

    def get_signal_statistics_for_display(self, df_result):
        """获取信号统计信息用于显示"""
        # 统计综合信号分布
        signal_counts = {
            '强烈买入': 0, '买入': 0, '持有': 0, 
            '卖出': 0, '强烈卖出': 0, '错误': 0, '数据不足': 0
        }
        
        for signal in df_result['综合信号']:
            if '强烈买入' in signal:
                signal_counts['强烈买入'] += 1
            elif '买入' in signal:
                signal_counts['买入'] += 1
            elif '持有' in signal:
                signal_counts['持有'] += 1
            elif '卖出' in signal:
                signal_counts['卖出'] += 1
            elif '强烈卖出' in signal:
                signal_counts['强烈卖出'] += 1
            elif '错误' in signal:
                signal_counts['错误'] += 1
            elif '数据不足' in signal:
                signal_counts['数据不足'] += 1
        
        return signal_counts

    def get_signal_statistics(self, df_result):
        """获取信号统计信息（保持原有功能，打印到控制台）"""
        print("\n=== 信号统计 ===")
        print(f"总股票数量: {len(df_result)}")
        
        signal_counts = self.get_signal_statistics_for_display(df_result)
        
        for signal_type, count in signal_counts.items():
            if count > 0:
                percentage = (count / len(df_result)) * 100
                print(f"{signal_type}: {count}只 ({percentage:.1f}%)")
        
        return signal_counts


# 使用示例
if __name__ == "__main__":
    strategy = ComprehensiveStrategy(
        lookback_period=10,
        volume_multiple=2.0,
        short_window=5,
        long_window=20,
        bollinger_window=20,
        num_std=2,
        mean_reversion_window=10,
        z_score_threshold=1.0
    )

    signals = strategy.generate_signals()
    
    # 打印前几行结果到控制台
    print("综合策略信号结果:")
    print(signals.head(15))
    
    # 获取信号统计
    strategy.get_signal_statistics(signals)
    
    # 可视化展示
    strategy.visualize_signals(signals)
    
    # 可选：保存到Excel
    # strategy.save_to_excel(signals)